/*
* Copyright (c) Microsoft Corporation.  All Rights Reserved.
*
* Handles interaction between the mocking framework and the detours framework
*/
#pragma once
#include "Mock10.h"
#include <windows.h>
#include "DetourFunction.h"
#include "DetourShim.h"

#pragma push_macro("SEALED")
#undef SEALED

#if _MSC_VER >= 1400
# define SEALED sealed
#else
# define SEALED
#endif

namespace WEX { namespace Common
{
    template <typename T>
    class PtrRefCount;
}}

namespace Mock10
{
    class Mock;
}

namespace Mock10 { namespace Private
{
    class Detour;

    template <typename TPrototype, typename TDetourPrototype>
    class DetourFunction;

    // Forward declaration to work around a compiler bug.  If there is a CriticalSectionLock class defined in a global scope,
    // the friend declaration below will apply to that class rather than to Mock10::Private::CriticalSectionLock
    class CriticalSectionLock; 

    // We need to implement our own CriticalSection class, so all calls to critical synchronization APIs are made from within Mock10.dll.
    // The header-only WEX::Common::CriticalSection will not work here since this header is published; so the code could end up being
    // compiled into another DLL.
    class MOCK10_API CriticalSection
    {
    friend class CriticalSectionLock;

    public:
        CriticalSection();
        ~CriticalSection();

    protected:
        void Enter();
        void Leave();

    private:
        CRITICAL_SECTION m_criticalSection;
    };

    class MOCK10_API CriticalSectionLock
    {
    public:
        explicit CriticalSectionLock(CriticalSection& criticalSection);
        ~CriticalSectionLock();

        void Unlock();

    private:
        CriticalSectionLock(const CriticalSectionLock&);
        CriticalSectionLock& operator=(const CriticalSectionLock&);

        void Lock();

        CriticalSection& m_criticalSection;
        bool m_locked;
    };

    // Forward declaration to work around a compiler bug.  If there is a ReadCriticalSectionLock or WriteCriticalSectionLock class defined in a global scope,
    // the friend declaration below will apply to that class rather than to Mock10::Private::ReadCriticalSectionLock or Mock10::Private::WriteCriticalSectionLock
    class ReadCriticalSectionLock;
    class WriteCriticalSectionLock;

    // We need to implement our own ReadWriteCriticalSection class, so all calls to critical synchronization APIs are made from within Mock10.dll
    class MOCK10_API ReadWriteCriticalSection
    {
    friend class ReadCriticalSectionLock;
    friend class WriteCriticalSectionLock;

    public:
        ReadWriteCriticalSection();
        ~ReadWriteCriticalSection();

        HRESULT GetStatus();

    private:
        bool WriteLock();
        void WriteUnlock();
        bool ReadLock();
        void ReadUnlock();

        CriticalSection m_readCriticalSection;
        HANDLE m_semaphore; // Using raw handle because Handle.h is not published
        volatile long m_writingThread; // Volatile ensures reads are always performed rather than reading from the cache
        int m_readLocks;
        HRESULT m_status;
    };

    // Allows reentrant read locks on the current thread if the thread already holds the read or write lock
    // Does not allow for promotion from read to write
    // Not safe to free out-of-order; objects must be destructed in the opposite order in which they were constructed
    class MOCK10_API ReadCriticalSectionLock
    {
    public:
        explicit ReadCriticalSectionLock(ReadWriteCriticalSection& lock);
        ~ReadCriticalSectionLock();

    private:
        ReadCriticalSectionLock(const ReadCriticalSectionLock&);
        ReadCriticalSectionLock& operator=(const ReadCriticalSectionLock&);

        ReadWriteCriticalSection& m_lock;
        bool m_locked;
    };

    // Allows rentrant write locks on the current thread if the thread already holds the write lock
    // Not safe to free out-of-order; objects must be destructed in the opposite order in which they were constructed
    class MOCK10_API WriteCriticalSectionLock
    {
    public:
        explicit WriteCriticalSectionLock(ReadWriteCriticalSection& lock);
        ~WriteCriticalSectionLock();

    private:
        WriteCriticalSectionLock(const WriteCriticalSectionLock&);
        WriteCriticalSectionLock& operator=(const WriteCriticalSectionLock&);
        
        ReadWriteCriticalSection& m_lock;
        bool m_locked;
    };

    class MOCK10_API DetourMap SEALED
    {
    friend class Detour;

    public:
        DetourMap();
        ~DetourMap();

    private:
        struct Impl;
        Impl* m_pImpl;
        HRESULT m_status;
    };

    class MOCK10_API DetourIndex SEALED
    {
    template <typename TPrototype, typename TDetourFunction>
    friend struct DetourShim;

    public:
        ~DetourIndex();
        int Next();

    private:
        DetourIndex(Detour* pDetour);

        struct Impl;
        Impl* m_pImpl;
        HRESULT m_status;
    };

    class MOCK10_API AutoDisableDetour SEALED
    {
    template <typename TPrototype, typename TDetourFunction>
    friend struct DetourShim;

    public:
        ~AutoDisableDetour();

    private:
        explicit AutoDisableDetour(Detour* pDetour);
        AutoDisableDetour& operator=(const AutoDisableDetour&); // Not implemented
        AutoDisableDetour(const AutoDisableDetour&); // Not implemented

        Detour* m_pDetour;
    };

    class MOCK10_API AutoDisableDetours SEALED
    {
    public:
        AutoDisableDetours();
        ~AutoDisableDetours();

    private:
        AutoDisableDetours& operator=(const AutoDisableDetours&); // Not implemented
        AutoDisableDetours(const AutoDisableDetours&); // Not implemented

        unsigned long m_lastError;
    };

    class MOCK10_API Detour
    {
    template <typename TPrototype, typename TDetourFunction>
    friend struct DetourShim;

    template <typename TPrototype, typename TDetourPrototype>
    friend class DetourManager;

    template <typename TPrototype, typename TDetourPrototype, typename TDerived>
    friend class DetourFunctionImpl;

    template <typename T>
    friend class WEX::Common::PtrRefCount;
    
    friend class AutoDisableDetour;
    friend class DetourThunk;
    friend class ExternalMockManager;
    friend class ::Mock10::Mock;

    friend MOCK10_API void* __stdcall ResolveFunctionPointer(void* pFunction);

    public:
        // Return the actual target function for the Detour instance
        // Note that this is also called the 'Trampoline' when the detour is in place.
        void* GetTarget() const;

        static bool __stdcall IsAddressInModule(void* pAddress);
        static bool __stdcall IsDisabled(Detour* pDetour);  // Whether or not the specified detour is disabled for this thread
        static bool __stdcall IsDisabled(); // Whether or not all detours are disabled for this thread

    protected:
        ~Detour();

    private:
        Detour(void* pTarget, void* pShim);

        struct Impl;
        Impl* m_pImpl;

        static MockStatus __stdcall Attach(void* pTarget, void* pDetourFunction, void* pShim, DetourMap*& pDetourMap, unsigned long& mapRefCount);
        static MockStatus __stdcall Detach(void* pTarget, void* pDetourFunction, DetourMap*& pDetourMap, unsigned long& mapRefCount);
        static ReadWriteCriticalSection& __stdcall GetGlobalCriticalSection();
        static bool __stdcall IsValid(void* pTarget, void* pDetour, DetourMap* const & pDetourMap);
        static void* CodeFromPointer(void* pTarget);

        unsigned long m_refCount;

        unsigned long AddRef();
        unsigned long Release();

        void* GetDetour(int index) const;
    };

    template <typename TPrototype, typename TDetourFunction>
    class DetourManager SEALED
    {
    template <typename TPrototype, typename TDetourFunction>
    friend struct DetourShim;

    public:
        static MockStatus __stdcall Attach(TPrototype* pTarget, TDetourFunction* pDetourFunction)
        {
            return Detour::Attach(pTarget, pDetourFunction, DetourShim<TPrototype, TDetourFunction>::ShimPointer(), s_pDetours, s_refCount);
        }

        static MockStatus __stdcall Detach(TPrototype* pTarget, TDetourFunction* pDetourFunction)
        {
            return Detour::Detach(pTarget, pDetourFunction, s_pDetours, s_refCount);
        }

        static bool __stdcall IsValid(TPrototype* pTarget, Detour* pDetour)
        {
            return Detour::IsValid(pTarget, pDetour, s_pDetours);
        }

    private:
        DetourManager();
        ~DetourManager();

        static DetourMap* s_pDetours;
        static unsigned long s_refCount;
    };

    template <typename TPrototype, typename TDetourFunction>
    DetourMap* DetourManager<TPrototype, TDetourFunction>::s_pDetours = NULL;

    template <typename TPrototype, typename TDetourFunction>
    unsigned long DetourManager<TPrototype, TDetourFunction>::s_refCount = 0;
} /* namespace Private */ } /* namespace Mock10 */

#pragma pop_macro("SEALED")